import pymysql

class MysqlHelper:
    def __init__(self, host, port, user, password, database):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.conn = None
        self.cursor = None

    def connect(self):
        try:
            self.conn = pymysql.connect(host=self.host, port=self.port, user=self.user, password=self.password, database=self.database)
            self.cursor = self.conn.cursor()
        except pymysql.MySQLError as e:
            print("数据库连接失败:", e)

    def close(self):
        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()

    def insert(self, sql, params):
        count = None
        try:
            self.connect()
            count = self.cursor.execute(sql, params)
            self.conn.commit()
        except pymysql.MySQLError as e:
            print("插入数据失败:", e)
            return 0
        return count
    
    def fetch_one(self, sql, params=None):
        result = None
        try:
            self.connect()
            self.cursor.execute(sql, params)
            result = self.cursor.fetchone()
        except pymysql.MySQLError as e:
            print("查询数据失败:", e)
        return result

class person:
    def register(self):
        username = input("请输入用户名：")
        password = input("请输入密码：")
        sql = "INSERT INTO tb_users (username, password) VALUES (%s, %s)"
        params = (username, password)
        helper = MysqlHelper('localhost', 3306, 'root', '123456', 'test')
        count = helper.insert(sql, params)
        if count > 0:
            print("注册成功")
        else:
            print("注册失败")
        helper.close()

    def login(self):
        username = input("请输入用户名：")
        password = input("请输入密码：")
        sql = "SELECT * FROM tb_users WHERE username = %s AND password = %s"
        params = (username, password)
        helper = MysqlHelper('localhost', 3306, 'root', '123456', 'test')
        result = helper.fetch_one(sql, params)
        if result:
            print("登录成功")
        else:
            print("登录失败")
        helper.close()

if __name__ == '__main__':
    p = person()
    while True:
        print("1.注册")
        print("2.登录")
        print("3.退出")
        choice = input("请输入操作：")
        if choice == '1':
            p.register()
        elif choice == '2':
            p.login()
        elif choice == '3':
            break
        else:
            print("输入错误，请重新输入")
        print()


